package HMM;

import TrainingData.TrainingInstance;
import TrainingData.TrainingInstanceCollection;

/**
 * @author Gergely Neu
 *
 * Maximum Margin training of an HMM.
 */
public class MaxMargin extends TrainingMethod {
	
	protected double _lossConstant;
	
	/**
	 * Constructor for MaxMargin. Extra parameter: lossConstant.
	 * 
	 * @param iterNr number of iterations to carry out.
	 * @param trainingData the data used for training. 
	 * @param testData the data used for testing.
	 * @param lossConstant the loss constant used to decrease scores of gold standard features.
	 */
	public MaxMargin(int iterNr, TrainingInstanceCollection trainingData,TrainingInstanceCollection testData, double lossConstant){
		super(iterNr,trainingData,testData);
		
		_lossConstant = lossConstant;
	}
	

	@Override
	protected Updates iterationStep(DoubleSparseMatrix transitionMx, DoubleSparseMatrix emissionMx) {
		
		int numStates = transitionMx.getColNr();
		int numEmissions = emissionMx.getColNr();
		
		DoubleSparseMatrix deltaTransMx = new DoubleSparseMatrix(numStates,numStates, 0);
		DoubleSparseMatrix deltaEmissMx = new DoubleSparseMatrix(numStates,numEmissions, 0);
		
		for (TrainingInstance instance : _trainingData) {
			
			int[] emissions = instance.getEmissions();
			int[] states = instance.getStates();
			
			int L = emissions.length;
			
			for (int i = 1; i < L; i++) {
				transitionMx.increment(states[i-1],states[i], - _lossConstant);
				
				emissionMx.increment(states[i], emissions[i], - _lossConstant);
			}
			
			emissionMx.increment(states[0], emissions[0], - _lossConstant);
			
			HMMCommonFunctions.computeViterbiScores(emissions, emissionMx, transitionMx);
			
			int[] decodedStateArray = HMMCommonFunctions.decodeFrom(emissions, L, 0);
			
			for (int i = 1; i < L; i++) { 
				deltaTransMx.increment(decodedStateArray[i-1], decodedStateArray[i], -1);
				
				deltaEmissMx.increment(decodedStateArray[i],emissions[i], -1);
				
				deltaTransMx.increment(states[i-1],states[i], 1);
				
				deltaEmissMx.increment(states[i],emissions[i], 1);
			}
			
			deltaEmissMx.increment(decodedStateArray[0],emissions[0], -1);
			
			deltaEmissMx.increment(states[0],emissions[0], 1);

			for (int i = 1; i < L; i++) {
				transitionMx.increment(states[i-1], states[i], _lossConstant);
				
				emissionMx.increment(states[i],emissions[i], _lossConstant);
			}
			
			emissionMx.increment(states[0], emissions[0], - _lossConstant);	
		}
		
		return new Updates(deltaTransMx, deltaEmissMx);
	}

}
